import logging
from collections import Counter
from tqdm import tqdm
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from pinecone import Pinecone, ServerlessSpec

# 配置logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')

# 初始化Pinecone客户端
pinecone = Pinecone(api_key="29788efd-4df4-41c3-a962-7937ed3558c7")

# 索引名称
index_name = "mnist2-index"

# 检查并创建索引
existing_indexes = pinecone.list_indexes()
if any(index['name'] == index_name for index in existing_indexes):
    logging.info(f"索引 '{index_name}' 已存在，正在删除...")
    pinecone.delete_index(index_name)
    logging.info(f"索引 '{index_name}' 已成功删除。")
else:
    logging.info(f"索引 '{index_name}' 不存在，将创建新索引。")

# 创建新索引
logging.info(f"正在创建新索引 '{index_name}'...")
pinecone.create_index(
    name=index_name,
    dimension=64,  # MNIST 每个图像展平后是一个 64 维向量
    metric="euclidean",  # 使用欧氏距离
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)
logging.info(f"索引 '{index_name}' 创建成功。")

# 连接到索引
index = pinecone.Index(index_name)
logging.info(f"已成功连接到索引 '{index_name}'。")

# 加载MNIST数据集
digits = load_digits(n_class=10)
X = digits.data
y = digits.target

# 划分数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 上传数据到Pinecone
vectors = []
for i in tqdm(range(len(X_train)), desc="上传数据到Pinecone"):
    vector_id = str(i)
    vector_values = X_train[i].tolist()
    metadata = {"label": int(y_train[i])}
    vectors.append((vector_id, vector_values, metadata))

batch_size = 1000
for i in tqdm(range(0, len(vectors), batch_size), desc="上传数据进度"):
    batch = vectors[i:i + batch_size]
    index.upsert(batch)

logging.info("成功上传了 1437条数据到 Pinecone")

# 测试准确率
query_data = X_test.tolist()
results = [index.query(vector=x, top_k=11, include_metadata=True) for x in tqdm(query_data, desc="测试k=11的准确率")]

# 计算准确率
correct_predictions = 0
for i, result in enumerate(results):
    true_label = y_test[i]
    label_counts = Counter([match['metadata']['label'] for match in result['matches']])

    # 检查label_counts是否为空
    if label_counts:
        predicted_label = label_counts.most_common(1)[0][0]
        if predicted_label == true_label:
            correct_predictions += 1
    else:
        # 如果label_counts为空，跳过这个样本
        print(f"警告: 查询结果为空，无法预测标签。样本索引: {i}")

# 计算准确率时只考虑那些查询结果非空的样本
accuracy = correct_predictions / len([result for result in results if result['matches']])
logging.info(f'当k=11时，使用Pinecone的准确率为: {accuracy:.2%}')